# -*-coding:utf-8-*-
"""
统计分析
"""

import json
import os

from data_full_view.common import datasets


def _group_by_names(dataframe, names):
    count_dict = {}
    for name in names:
        d = dataframe.groupby(name).size().to_dict()
        count_dict[name] = d
    return count_dict


def count_train_data_category(data_path, outf_dir=None):
    """
    统计训练数据每个类别的数量
    :return:
    """
    result_file_name = 'category_statistical.json'
    keys = datasets.get_keys()
    data_set = datasets.Datasets(data_path)
    train_data = data_set.train_data()
    names = keys[2:]
    category_count = _group_by_names(train_data, names)
    if outf_dir:
        f = open(os.path.join(outf_dir, result_file_name), 'w')
        json.dump(category_count, f, indent=4)


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(description='统计分析')
    parser.add_argument('--data_path', required=True, help="下载文件的解压后地址")
    parser.add_argument('--out_dir', help='输出文件位置' )
    args = parser.parse_args()
    data_path = args.data_path
    out_dir = args.out_dir
    count_train_data_category(data_path=data_path, outf_dir=out_dir)
